from mpl_toolkits.mplot3d import Axes3D
import numpy as np
from pylab import *
import matplotlib.pyplot as plt
import time as timing #Renamed because time variable is already in use. Careful!
import math
import copy #Used to make a copy of grid objects.
from multiprocessing.queues import SimpleQueue
from multiprocessing import Process

#New instance either is all zeroes, or is copied off of an existing grid instance (template).
#Largely used as a storage unit for the displacement of the simulated sheet at different points.
#However, is also used to store mechanical properties of the sheet as values (this allows mechanical
#parameters to vary across the sheet, such as elastic modulus and density).
#In all honesty, this largely just creates functions to request pieces of the array
#in plain english. Grid.grid[:,0,1] vs Grid.getcolumn(0)
#Could be cleaned up, and sometimes Sheet just directly modifies the underlying
#numpy array.
#The third dimension to the matrix, tsize, is used represent the various timesteps
# 0 = past, 1 = present, 2 = future.
class Grid:
	def __init__(self, xsize = 5, ysize = 5, tsize = 3, template = None):
		if template == None:
			self.grid = np.zeros((xsize, ysize, tsize))
			self.t = tsize
			self.xsize = xsize
			self.ysize = ysize
		else:
			self.grid = np.copy(template.grid)
			self.t = template.t		

	def getval(self, x, y, time=1):
		if self.t == 1:
			return self.grid[x,y]
		else:
			return self.grid[x,y,time]

	def setval(self, x, y, val, time=1):
		if self.t == 1:
			self.grid[x,y] = val
		else:
			self.grid[x,y, time] = val

	def getrow(self, column, time=1):
		if self.t == 1:
			return self.grid[:,column]
		else:
			return self.grid[:,column, time]

	def getcolumn(self, row, time=1):
		if self.t == 1:
			return self.grid[row,:]
		else:
			return self.grid[row,:, time]

	def setcolumn(self, row, val, time=1):
		if self.t == 1:
			self.grid[row,:] = val
		else:
			self.grid[row,:, time] = val
		
	def setall(self, val):
		for x in range(self.xsize):
			for y in range(self.ysize):
				if self.t == 1:
					self.grid[x,y] = val
				else:
					for t in range(self.t):
						self.grid[x,y,t] = val	
			
	def shape(self):
		return self.grid.shape

	def multiply(self, val):
		self.grid = self.grid * val

	#Can directly set the grid class to match a numpy array.
	def importnparray(self, nparray):
		self.grid = np.copy(nparray)

#Module for all the derivative calculations.
class MathModule:
	#Based off a recommendation of Paul Reimer.
	#Uses the symmetry assumption that near end points, a point distance x away from the
	#point in question should have the same value as the point a distance -x from the point.
	#Also, assumes a fixed end (displacement = 0 at boundary).
	#Normally gathers a series of points necessary to calculate the derivative.
	def getpoints(self, point, distance, array, fixed = False):
		vals = []
		index = []
		start = point - distance
		end = point + distance + 1
		points = range(start, end)
		for val in points:
			if abs(val) < (array.size - 1):
				vals.append(val)
			else:
				if fixed:
					vals.append(None)
				else:
					vals.append((array.size - 1) - (val - (array.size - 1)))
		for a in vals:
			if a == None:
				index.append(0.0)
			else:
				index.append(array[a])			
		return index
	
	def secondderivative(self, array, point, stepsize):
		vals = self.getpoints(point, 1, array, fixed = False)		
		if self.equalvals(vals):
			return 0.0
		derivative = (vals[2] - 2*vals[1] + vals[0]) / stepsize**2
		return derivative

	#High accuracy simply makes use of a greater number of points in it's calculation
	#This does not always translate to greater accuracy of the result if the points
	#are spaced too far away from the point being calculated.
	def secondderivativehighaccuracy(self, array, point, stepsize):
		vals = self.getpoints(point, 2, array)
		if self.equalvals(vals):
			return 0.0

		derivative = (-vals[4] + 16*vals[3] - 30*vals[2] + 16*vals[1] - vals[0]) / 12*stepsize**2
		return derivative

	def fourthderivative(self, array, point, stepsize):
		vals = self.getpoints(point, 2, array, fixed = False)
		if self.equalvals(vals):
			return 0.0
		derivative = (vals[4] - 4*vals[3] + 6*vals[2] - 4*vals[1] + vals[0]) / stepsize**4	
		return derivative

	def fourthderivativehighaccuracy(self, array, point, stepsize):
		vals = self.getpoints(point, 3, array)
		if self.equalvals(vals):
			return 0.0
		derivative = (-vals[6] + 12*vals[5] - 39*vlas[4] + 56*vals[3] - 39*vals[2] + 12*vals[1] - vals[0]) / 6*stepsize**4
		return derivative

	#Used so that values less than 0.0001 are simply returned as zero (to prevent
	#values like 3*10^-14 being returned, makes it more readable.
	def equalvals(self, vals):
		check = None
		for i in vals:
			if check == None:
				check = i
			else:
				if abs(check - i) >= 0.0001:
					return False
		return True


#Represents the actual oscillating 2D sheet.
class Sheet:
	#Sets all internal variables for simulation.
	#The spacing variables represent the distance in time or space a step along the grid represents.
	#Keep in mind that the size of the grid times the grid spacing represents the total area actually simulated.
	def __init__(self, xsize = 20, ysize = 20, timespacing=0.05, distspacing=0.5):
		self.xsize = xsize
		self.ysize = ysize
		self.xspacing = distspacing
		self.tspacing = timespacing
		self.grid = Grid(self.xsize, self.ysize)
		self.math = MathModule()
		self.setmaterialparameters()

	#Sets all the 'physical' parameters, values that are mechanical properties of the sheet material.
	#User can either manually set values or program uses defaults.
	#Users can also 'unpickle' a numpy array for the emod and density, allowing for more complex
	#changes in mechanical parameters across the sheet.
	def setmaterialparameters(self):
		#Uses the pickle module to unpack a stored numpy array.	
		def unpack():
			print('Unpickles a numpy array.')
			check = False #Does not allow a user to pass until the arrays passed through
				      #pickle are the the same dimensions as the grid.
			while check == False:
				import pickle
				print('Make sure the arrays are the same shape as sheet.')
				emodpickle = raw_input("Name of pickled elastic modulus?(default:emod.pkl): ")
				densitypickle = raw_input("Name of pickled density?(default:density.pkl): ")
				if emodpickle == '': #Sets a default name if none is provided.
					emodpickle = 'emod.pkl'
				if densitypickle == '':
					densitypickle = 'density.pkl'
				emodfile = open(emodpickle, 'rb') #First loads the actual file.
				densityfile = open(densitypickle, 'rb')
				emodarray = pickle.load(emodfile) #Then is unpacked by the pickle module.
				densityarray = pickle.load(densityfile)
				print('grid shape: ' + str(self.grid.shape()[0:2]))
				if emodarray.shape != self.grid.shape()[0:2]:
					print('emod shape: ' + str(emodarray.shape))
				if densityarray.shape != self.grid.shape()[0:2]:
					print('density shape: ' + str(densityarray.shape))
				else:
					check = True 
			self.elasticmodulus.importnparray(emodarray) #If the array is the right size, it sets
			self.density.importnparray(densityarray)	#the array as the physical parameters.
	
		self.density = Grid(self.xsize, self.ysize, 1)
		self.elasticmodulus = Grid(self.xsize, self.ysize, 1)
		custom  = raw_input('Unpickle elastic modulus and density? (y/n): ')
		if custom == 'y':
			unpack()
		else:		
			modifymaterialparameters = bool(raw_input("Specify material parameters? (0 or 1): "))
			if not modifymaterialparameters: #Sets a series of default parameters for the whole sheet.
				emod = 18.0	
				density = 4.0
				self.poisson = 0.37
			else:
				emod = float(raw_input("Elastic modulus: ")) #...or a user defined parameter.
				density = float(raw_input("Density: " ))

			self.elasticmodulus.setall(emod)
			self.density.setall(density)

		self.poisson = raw_input("Poisson ratio(default: 0.37): ")
		if self.poisson == '':
			self.poisson = 0.37
		else:
			self.poisson = float(self.poisson)
		self.D = Grid(1,1,1,self.elasticmodulus)
		self.thickness=0.5 #Actual thickness of the sheet.
		const = self.thickness**3/(12*(1-self.poisson**2))
		self.D.multiply(const) #Because D is dependent on the elastic modulus, it too can vary across the sheet.

	#Interates through each timestep, rendering the sheet as it progresses.
	#Is not run automatically, must be run as a separate step from the constructor.
	#Drivingfreq, maxamplitude, and setdrive are all used by the drive() function.
	#Damping is used by the damping() function
	def mainloop(self, drivingfreq, endtime, maxamplitude, threads, setdrive='mid', damping=-0.2):
		time = 0
		self.notified = False #See method checkstability().
		self.fixed = False
		self.amplitude = maxamplitude #Maximum amplitude of the sheet driver.
		self.renderscene() #Sets up the initial rendering of the sheet.
		self.threadnumber = threads #Number of child threads created to perform calculations.
		self.drivetype = setdrive
		self.dampingcoeff = damping
		while time < endtime:
			try:
				self.drive(drivingfreq, self.tspacing, time, maxamplitude)
				self.thread()
				self.grid.grid[:,:,0] = np.copy(self.grid.grid[:,:,1])
				self.grid.grid[:,:,1] = np.copy(self.grid.grid[:,:,2])
				time += 1.0
				#if time % 5 == 0:	
				self.visualize2D(time) #Renders the sheet every step (or multiple steps)	
				print(int(time))	
				if self.notified == True: #If checkstability() is tripped, ends the calculations (unstable).
					break
			except KeyboardInterrupt:
				print('\n Interrupted, stopping')
				break

		response = raw_input('Close? (Hit Enter) ')	
		plt.close('all') #Closes any render windows still open.

	#Creates the inital scene. See matplotlib documentation online if you want a deeper understanding of how
	#this plot is rendered.
	def renderscene(self):
		self.X = np.arange(0, self.xsize)
		self.Y = np.arange(0, self.ysize)
		self.X, self.Y = np.meshgrid(self.X, self.Y)
		self.Z = self.grid.getval(self.X, self.Y)
		self.fig = plt.figure(figsize = (14, 11), dpi = 80) #Specifies the actual size of the sheet render window.
		self.ax = self.fig.gca(projection='3d')
		self.surf = self.ax.plot_surface(self.X, self.Y, self.Z, rstride=1, cstride=1, cmap=cm.coolwarm, linewidth=0, antialiased=False)

	#Pause/render method for visualizing wave propogation.
	def visualize2D(self, settime, done=False):
		self.surf.remove()
		self.Z = self.grid.getval(self.X, self.Y)
		self.surf = self.ax.plot_surface(self.X, self.Y, self.Z, rstride=1, cstride=1, cmap=cm.jet, vmin = -self.amplitude*1.1, vmax = self.amplitude*1.1, linewidth=0, antialiased=True)
		self.ax.set_zlim((-int(self.amplitude*4.0),int(self.amplitude*4.0)))
		if settime % 2 == 0:
			plt.draw()
		if not done:
			val = int(settime)
			if val < 10:
				val = '000' + str(val)
			elif val < 100:
				val = '00' + str(val)
			elif val < 1000:
				val = '0' + str(val)
			elif val < 10000:
				val = str(val)
			fname = 'image' + str(self.drivetype) + str(val) + '.png'
			#Uncomment the next line if making an animation. 			
			#matplotlib.pyplot.savefig(fname, edgecolo='w', format='png')
		timing.sleep(0.01)		


	#Creates the effect of a driving force on the sheet, either at a point or in a line. 
	def drive(self, drivingfreq, timespacing, curtime, amplitude):
		sin, pi = np.sin, np.pi
		oscillation = amplitude * sin(2 * np.pi * drivingfreq * timespacing * curtime)
		#Set pulse to true if you only want a wave packet set on the sheet.
		pulse = False
		if pulse = True:
			if curtime > 10 and abs(oscillation) < 0.05 and self.fixed == False:
				oscillation = 0
				self.fixed = True
				print('Fixed')
			if self.fixed:
				oscillation = 0
			oscillation = amplitude
			self.grid.setcolumn(0, oscillation)
		if self.drivetype == 'mid':
			self.grid.setval(self.xsize/2, self.ysize/2,oscillation)
		if self.drivetype == 'side':
			self.grid.setcolumn(0, oscillation)
		if self.drivetype == 'end':
			self.grid.setval(self.xsize/2, 0, oscillation)

	#Checks if a result value on sheet exceeds the max driving frequency amplitude.
	#If the warning is thrown, this probably means the amplitude is growing  unbounded (unrealistic).
	def checkstability(self, val, amplitude, x, y):
		if abs(val) >= 10.0 * amplitude and not self.notified:
			print('Instability!')
			self.notified = True
	
	#Multithreading!
	#Breaks down the sheet into even chunks along the x axis for the child processes.
	#Then reassembles the grid based on the returns from the children.
	#Can largely be ignored with regards to how the actual calculations are performed,
	#this is only to divy up the responsibilities among the child processes.
	#Yes, currently child process are created/destroyed each iteration. Currently working
	#on a better solution in model2Dtest.py.
	def thread(self):
		threadarray = []
		q = SimpleQueue()
		start = 0
		sizeperthread = self.xsize / self.threadnumber
		childID = 0
		for x in range(self.threadnumber):
			newchild = Process(target=self.shake(q,childID,[start,start+sizeperthread]))
			threadarray.append(newchild)
			threadarray[x].start()
			start += sizeperthread
			childID += 1

		gridarray = []
		returnedthreads = 0
		while returnedthreads != self.threadnumber:
			if not q.empty():
				val = q.get()
				gridarray.append(val)
				returnedthreads += 1

		increment = -1
		val = 0	
		newarray = None
		while increment != self.threadnumber:
			for x in gridarray:
				if x[0] == increment + 1:
					if newarray == None:
						newarray = np.copy(x[1])
					else:
						newarray = np.concatenate((newarray, x[1]))
			increment += 1
		
		self.grid.grid[:,:,2] = np.copy(newarray)
		for x in threadarray:
			x.terminate()		
		return 0
		
		
	#Called function by thread().
	#Used to actually simulate the shaking of the sheet.
	#Because child processes inherit memory space, we don't need args!...besides the return queue, child ID,
	#and grid portion that is the child's responsibility.
	#Puts the calculated portion of the grid into the queue then returns.
	def shake(self,queue,childID,portion):
		for x in range(portion[0],(portion[1])):
			column = self.grid.getcolumn(x)
			dx2column = self.constructdx2(x)
			for y in range(self.ysize):
				row = self.grid.getrow(y)
				increment = self.increment(x, y, dx2column, row, column)
				self.checkstability(increment, self.amplitude, x, y)
				self.grid.setval(x, y, increment, 2)

		gridportion = np.copy(self.grid.grid[portion[0]:portion[1],:,2])
		queue.put([childID, gridportion])
		return 0

	#Finds d2w/dx2 at each point along a column of y values, for use to find d4w/dx2dy2.
	def constructdx2(self, x):
		dwdx2column = np.zeros((self.ysize))
		for y in range(self.ysize):
			row = self.grid.getrow(y)
			dwdx2 = self.math.secondderivative(row, x, self.xspacing) 
			dwdx2column[y] =  dwdx2
		return dwdx2column	

	#Individual point change calculation, applied at every point on grid.
	def increment(self, x, y, dx2array, row, column):
		dwdy4 = self.math.fourthderivative(column, y, self.xspacing)
		dwdx4 = self.math.fourthderivative(row, x, self.xspacing)				
		dwdx2dy2 = self.math.secondderivative(dx2array[:], y, self.xspacing)
		#See pg. 22 in the lab notebook for this calculation in written form.
		seg1 = self.D.getval(x,y,0)*(dwdy4 + dwdx4 + dwdx2dy2)
		seg2 =  self.getdensity(x,y) * self.thickness*(-2*self.grid.getval(x,y)+self.grid.getval(x,y,0))/self.tspacing**2
		seg3 =  self.damping(x,y)*-self.grid.getval(x,y,0)/(2*self.tspacing)
		seg4 = 	self.getdensity(x,y)*self.thickness/self.tspacing**2+self.damping(x,y)/(2*self.tspacing)
		increment = -(seg1 + seg2 + seg3)/seg4
		self.checkstability(increment, self.amplitude, x, y)
		return increment	

	#Calculates the damping for a position on the sheet (should be greatest near the edge).
	#The damping coeff increases linearly up to some highest value.
	def damping(self, x, y):
		coeff = 0.0
		damping_thickness_x = 0.15
		damping_thickness_y = 0.05
		if y <= self.ysize * damping_thickness_y:
			coeff = self.dampingcoeff - self.dampingcoeff / (self.ysize * damping_thickness_y) * y
		if y >= self.ysize * (1.0 - damping_thickness_y):
			coeff = self.dampingcoeff / (self.ysize * damping_thickness_y) * (y - (1 - damping_thickness_y)*self.ysize)	
		if x >= self.xsize * (1.0 - damping_thickness_x):
			coeff = self.dampingcoeff / (self.xsize * damping_thickness_x) * (x - (1 - damping_thickness_x)*self.xsize)
		return coeff

	def getdensity(self, x, y):
		return self.density.getval(x,y)

#Quickly sets the components necessary for a basic simulation.
def sim(xsize=15, ysize = 15, time=20, drive='side'):
	threads = thread(xsize)
	tspacing = 0.05 
	xspacing = 0.5
	matplotlib.interactive(True)
	drivingamplitude = 0.5
	drivingfreq = 0.25
	damping = 0.2
	sheet = Sheet(xsize, ysize, tspacing, xspacing)
	sheet.mainloop(drivingfreq, time, drivingamplitude, threads, setdrive=drive, damping=damping)	 

#The number of threads should be a divisor of the size in the x direction, so that the job can be split
#evenly among the child processes.
def thread(xsize):
	extra = 1
	while extra != 0:
		threads = int(raw_input('Number of threads? (must be a divisor of ' + str(xsize) + '): '))
		extra = xsize % threads
	return threads 
